"""LangGraph节点定义"""

from typing import Dict, Any
from zhipuai import ZhipuAI
import streamlit as st
from config import Config
from agent.utils import TextProcessor
from agent.utils import PDFProcessor
from prompts.summary_prompts import PAPER_SUMMARY_PROMPT
from prompts.mindmap_prompts import MINDMAP_PROMPT
import json
import re

class PaperPalNodes:
    """PaperPal Agent的节点定义"""
    
    def __init__(self, api_key: str):
        self.client = ZhipuAI(api_key=api_key)
        self.text_processor = TextProcessor()
        self.pdf_processor = PDFProcessor()
    
    def call_glm(self, prompt: str, max_tokens: int = 4000) -> str:
        """调用GLM-4-Flash模型"""
        try:
            response = self.client.chat.completions.create(
                model=Config.MODEL_NAME,
                messages=[
                    {"role": "user", "content": prompt}
                ],
                max_tokens=max_tokens,
                temperature=Config.TEMPERATURE,
                stream=False
            )
            return response.choices[0].message.content
        except Exception as e:
            st.error(f"GLM API调用失败: {str(e)}")
            return f"API调用失败: {str(e)}"

    def preprocess_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """文本预处理节点"""
        st.info("🔄 正在预处理文档...")

        try:
            raw_data = state.get("raw_text", "")  # 改为支持 dict 输入
            if not raw_data:
                st.error("❌ 未找到文档内容")
                state.update({
                    "preprocessing_complete": False,
                    "error": "No raw text found"
                })
                return state

            # 支持新版结构化输入：{"text": ..., "sections": [...]}
            if isinstance(raw_data, dict) and "text" in raw_data:
                raw_text = raw_data["text"]
                sections = raw_data.get("sections", [])
            elif isinstance(raw_data, str):
                raw_text = raw_data
                sections = []
            else:
                raise ValueError("raw_text 类型不支持，应为字符串或包含 'text' 键的字典")

            # 文本清理
            cleaned_text = self.text_processor.clean_text(raw_text)

            # 如果 sections 为空，再尝试章节划分（兼容未预分好章节的情况）
            if not sections:
                sections = self.text_processor.split_into_sections(cleaned_text)

            # 文本分块（处理长文档）
            text_chunks = self.text_processor.chunk_text(cleaned_text, max_tokens=8000)

            # 更新状态
            state.update({
                "cleaned_text": cleaned_text,
                "sections": sections,
                # "formulas": formulas,
                "text_chunks": text_chunks,
                "preprocessing_complete": True
            })

            st.success(f"✅ 预处理完成！")

        except Exception as e:
            st.error(f"❌ 预处理过程中发生错误：{str(e)}")
            st.write(f"错误详情：{type(e).__name__}: {str(e)}")
            state.update({
                "preprocessing_complete": False,
                "error": str(e)
            })

        return state

    def analysis_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """内容分析节点"""
        st.info("🧠 正在分析论文内容...")
        
        try:
            text_chunks = state.get("text_chunks", [])
            sections = state.get("sections", [])
            
            # 确保sections是字典列表格式
            if not sections:
                st.warning("⚠️ 未找到章节信息，跳过分析步骤")
                state.update({
                    "section_analyses": [],
                    "analysis_complete": True
                })
                return state
            
            # 分析每个章节
            section_analyses = []
            for i, section in enumerate(sections):
                try:
                    # 确保section是字典格式
                    if isinstance(section, dict):
                        section_title = section.get("title", f"Section {i+1}")
                        section_content = section.get("content", "")
                    else:
                        # 如果section是字符串，转换为字典格式
                        section_title = f"Section {i+1}"
                        section_content = str(section)
                    
                    if section_content.strip():
                        # 限制内容长度以避免超出API限制
                        content_preview = section_content[:2000] if len(section_content) > 2000 else section_content
                        
                        prompt = f"""请分析以下论文章节的主要内容：

    章节标题：{section_title}

    章节内容：
    {content_preview}

    请提供简洁的分析，包括：
    1. 主要观点
    2. 关键信息
    3. 重要结论

    分析结果："""
                        
                        analysis = self.call_glm(prompt, max_tokens=1000)
                        
                        section_analyses.append({
                            "title": section_title,
                            "analysis": analysis,
                            "content_length": len(section_content)
                        })
                        
                except Exception as e:
                    st.warning(f"⚠️ 分析章节 {i+1} 时出错：{str(e)}")
                    continue
            
            # 更新状态
            state.update({
                "section_analyses": section_analyses,
                "analysis_complete": True
            })
            
        except Exception as e:
            st.error(f"❌ 分析过程中发生错误：{str(e)}")
            state.update({
                "section_analyses": [],
                "analysis_complete": False,
                "error": str(e)
            })
        
        return state
    
    def summary_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """摘要生成节点"""
        st.info("📝 正在生成论文摘要...")
        
        # 获取前几个文本块用于摘要生成
        text_chunks = state.get("text_chunks", [])
        summary_text = "\n".join(text_chunks[:2])  # 使用前两个块
        
        # 生成摘要
        prompt = PAPER_SUMMARY_PROMPT.format(paper_content=summary_text)
        summary = self.call_glm(prompt, max_tokens=2000)
        
        # 更新状态
        state.update({
            "summary": summary,
            "summary_complete": True
        })
        
        st.success("✅ 摘要生成完成！")
        return state
    
    def mindmap_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """思维导图生成节点"""
        st.info("🗺️ 正在生成思维导图...")
        
        # 使用摘要或清理后的文本生成思维导图
        summary = state.get("summary", "")
        cleaned_text = state.get("cleaned_text", "")
        
        # 优先使用摘要，如果没有则使用原文的一部分
        source_text = summary if summary else cleaned_text[:3000]
        
        prompt = MINDMAP_PROMPT.format(paper_content=source_text)
        mindmap_code = self.call_glm(prompt, max_tokens=1500)
        
        # 提取Mermaid代码
        if "```mermaid" in mindmap_code:
            mindmap_code = mindmap_code.split("```mermaid")[1].split("```")[0].strip()
        elif "mindmap" in mindmap_code:
            # 如果没有代码块格式，尝试提取mindmap部分
            lines = mindmap_code.split('\n')
            mindmap_lines = []
            in_mindmap = False
            for line in lines:
                if "mindmap" in line:
                    in_mindmap = True
                    mindmap_lines.append(line)
                elif in_mindmap and line.strip():
                    mindmap_lines.append(line)
                elif in_mindmap and not line.strip():
                    break
            mindmap_code = '\n'.join(mindmap_lines)
        
        # 更新状态
        state.update({
            "mindmap": mindmap_code,
            "mindmap_complete": True
        })
        
        st.success("✅ 思维导图生成完成！")
        return state
    
    def is_valid_formula(self, formula_str: str) -> bool:
        """
        判断 LaTeX 公式是否为有效数学表达式，排除伪公式（如纯变量名、常数、数字幂等）。
        """
        # 去掉包裹的 $$，并标准化空格和大小写
        cleaned = formula_str.strip().strip("$").replace(" ", "").lower()

        # 空字符串或纯数字
        if not cleaned:
            return False
        if re.fullmatch(r'[0-9.eE+-]+[%°]?', cleaned):
            return False

        # 纯文本内容
        if re.fullmatch(r'\\text\{[^}]*\}', cleaned):
            return False

        # 数字的幂表达式，如 10^{11}
        if re.fullmatch(r'[0-9]+(\.[0-9]+)?\^\{?[0-9.]+\}?', cleaned):
            return False

        # 单一变量（允许下标、上标，但不能含运算符）
        if re.fullmatch(r'[a-z](?:_\{?[a-z0-9]+\}?|\^\{?[a-z0-9]+\}?)*', cleaned):
            return False

        # 太短的内容（容易是误识别）
        if len(cleaned) < 5:
            return False

        # 没有任何数学结构关键词
        math_keywords = [
            '=', '<', '>', '\\frac', '\\int', '\\sum', '\\lim',
            '\\log', '\\sin', '\\cos', '\\tan', '\\sqrt',
            '\\cdot', '\\begin{bmatrix}', '^'
        ]
        if not any(kw in cleaned for kw in math_keywords):
            return False

        # 纯数字和算术符号，不包含变量（如 "2+3*5"）
        if re.fullmatch(r'[\d\+\-\*/\^\.\(\)\{\}]+', cleaned):
            return False

        return True

    def formula_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        st.info("🧮 正在提取并解释数学公式...")

        sections = state.get("sections", [])
        if not sections:
            st.warning("⚠️ 未找到任何章节内容，无法提取公式")
            state.update({
                "formulas": [],
                "formula_explanation": [],
                "formula_complete": True
            })
            return state

        all_formulas = []
        formula_explanations = []

        base_prompt_template = """
        你是**数学公式专家**。请**严格判断**以下文本中是否存在**真正的、有数学含义的数学公式**。  
        只有当**内容包含明确的数学符号结构**（如等式、不等式、分数、根号、积分、求和、矩阵等），  
        并且**表达变量之间的数学关系时，才认定为有效数学公式**。  
        **请勿将简单的符号说明、缩写定义、变量名称或关键词当作数学公式**。

        请将每个识别出的数学公式转换为**标准 LaTeX 语法**，并用 `$$...$$` 包裹，以便于 Markdown 渲染；  
        同时，用中文**详细解释**每个公式的数学含义及用途，并尽可能用**通俗的语言**让毫无数学背景的人也能理解。
        说明变量代表什么，解释这个公式“想要表达什么”、“是做什么的” 。

        如果章节中没有有效数学公式，请返回空数组 `[]`，切勿捏造不存在的公式。  

        【章节内容】：  
        {content}

        👉 输出格式（严格 JSON 数组，格式如下）：  
        [
        {{
            "formula": "$$LaTeX公式$$", 
            "explanation": "中文详细解释，涉及公式的部分用$$...$$包裹，适当进行分段以提高可读性"
        }}
        ]

        ⚠️ 注意事项：  
        - 公式必须用 `$$` 包裹，且为标准 LaTeX 语法；  
        - LaTeX 中的反斜杠请写成 `\\\\`，以保证 JSON 解析正确；  
        - 严格只输出上述 JSON 内容，勿包含任何无关文本或多余标点。  
        """
        # 逐章节提取公式
        for sec_idx, section in enumerate(sections):
            content = section.get("content", "").strip()
            if not content:
                continue

            prompt = base_prompt_template.format(content=content)

            try:
                glm_response = self.call_glm(prompt, max_tokens=1500)
                match = re.search(r'\[.*\]', glm_response, re.DOTALL)
                if match:
                    json_str = match.group(0)
                    json_str_safe = re.sub(r'(?<!\\)\\(?![\\"])', r'\\\\', json_str)
                    formulas = json.loads(json_str_safe)
                    formulas = [f for f in formulas if self.is_valid_formula(f["formula"])]
                else:
                    formulas = []
            except Exception as e:
                st.warning(f"⚠️ 章节 {sec_idx+1} 公式提取失败: {e}")
                formulas = []

            if formulas:
                all_formulas.extend(formulas)
                for f in formulas:
                    formula_explanations.append({
                        "latex": f["formula"].strip("$"),
                        "context": f.get("explanation", "")
                    })

        # 如果没有任何公式，做数学概念层面的分析
        if not all_formulas:
            st.info("🧮 未检测到数学公式，开始进行数学概念层面分析...")

            concept_prompt = """
            你是数学专家。请根据以下文本，提取并解释其中的数学概念（如定义、定理、性质、核心思想等）。  
            请返回一个 JSON 数组，每条元素格式为：  
            [{{"latex": "$$\\text{{数学概念名}}$$", "explanation": "中文详细解释"}}]

            【文本内容】：  
            {content}

            请严格返回 JSON 数组，且不要输出其他内容。
            """

            full_text = "\n\n".join([sec.get("content", "") for sec in sections])
            prompt = concept_prompt.format(content=full_text)

            try:
                glm_response = self.call_glm(prompt, max_tokens=1500)
                match = re.search(r'\[.*\]', glm_response, re.DOTALL)
                if match:
                    json_str = match.group(0)
                    concepts = json.loads(json_str)
                else:
                    concepts = []
            except Exception as e:
                st.warning(f"⚠️ 数学概念提取失败: {e}")
                concepts = []

            # 转换为前端展示格式，latex字段空，context填解释
            formula_explanations = []
            for c in concepts:
                formula_explanations.append({
                    "latex": "",
                    "context": f"{c.get('concept', '')}：{c.get('explanation', '')}"
                })
        else:
            st.success("✅ 公式提取完成！")

        state.update({
            "formulas": all_formulas,
            "formula_explanation": formula_explanations,
            "formula_complete": True
        })

        return state
    
    def integration_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """结果整合节点"""
        st.info("🔗 正在整合分析结果...")
        
        # 检查所有组件是否完成
        components_status = {
            "预处理": state.get("preprocessing_complete", False),
            "内容分析": state.get("analysis_complete", False),
            "摘要生成": state.get("summary_complete", False),
            "思维导图": state.get("mindmap_complete", False),
            "公式解释": state.get("formula_complete", False),
            "论文推荐": state.get("recommendation_complete", False),
            "专家分析": state.get("expert_review_complete", False),
            "论文溯源": state.get("trace_complete", False)
        }
        
        # 生成处理报告
        completed_count = sum(components_status.values())
        total_count = len(components_status)
        
        # 更新状态
        state.update({
            "integration_complete": True,
            "all_complete": completed_count == total_count
        })
        
        if state["all_complete"]:
            st.success("🎉 所有分析任务完成！")
        else:
            st.warning(f"⚠️ 部分任务未完成 ({completed_count}/{total_count})")
        
        return state
    
    def recommendation_analysis_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """推荐分析节点：分析论文内容，提取关键词和主题"""
        st.info("🔍 正在分析论文以生成推荐...")

        try:
            # 提取摘要或使用摘要
            paper_summary = state.get("summary", "")
            if not paper_summary:
                paper_summary = state.get("cleaned_text", "")[:2000]

            # 调用模型提取关键词和主题
            prompt = f"""请分析以下论文内容，提取关键词和主题，用于生成相关论文推荐：

            论文摘要：
            {paper_summary}

            请提供：
            - 5-8个关键词
            - 主题领域（如：机器学习、自然语言处理、计算机视觉等）
            - 相关研究方向
            """
            response = self.call_glm(prompt, max_tokens=500)

            # 解析模型输出
            keywords = []
            themes = []
            directions = []

            if "关键词：" in response:
                keywords = response.split("关键词：")[-1].split("主题领域：")[0].strip().split(", ")
            if "主题领域：" in response:
                themes = response.split("主题领域：")[-1].split("相关研究方向：")[0].strip().split(", ")
            if "相关研究方向：" in response:
                directions = response.split("相关研究方向：")[-1].strip().split(", ")

            # 更新状态
            state.update({
                "recommendation_keywords": keywords,
                "recommendation_themes": themes,
                "recommendation_directions": directions,
                "recommendation_analysis_complete": True
            })

            st.success("✅ 推荐分析完成！")

        except Exception as e:
            st.error(f"❌ 推荐分析失败：{str(e)}")
            state.update({
                "recommendation_analysis_complete": False,
                "error": str(e)
            })

        return state
    
    def expert_review_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """模拟领域专家的评价"""
        st.info("📝 专家视角分析...")
        paper_content = state.get("summary", "")
        
        # 使用大模型API生成专家评价
        prompt = f"""
        假设你是该领域的专家，请对以下研究内容进行专业评审。请严格按照以下要求提供结构化反馈：
        
        ## 评审维度
        1. 创新性评估 (0-10分):
           - 评分：
           - 理由：
           - 与当前领域前沿的对比：

        2. 方法论评价：
           - 实验设计合理性 ("合理性强/中/弱", "主要缺陷"):
           - 数据分析（"数据充分性","分析方法适当性","可重复性")

        3. 理论贡献：
           - 理论深度：
           - 文献综述完整性：

        4. 实用价值：
           - 应用场景("列举1-3个")：
           - 技术转化可能性：

        5. 写作与表达：
           - 逻辑连贯性("高/中/低","理由")：
           - 图表有效性("高/中/低","理由")：
           - 专业术语准确性("高/中/低","理由")：

        ## 综合结论
        - 总体评分（1-5星）：
        - 推荐意见（是否推荐阅读，结合上述所有内容做总结）：
        - 对后续研究的建议：
        
        论文内容：
        {paper_content}
        """
        expert_feedback = self.call_glm(prompt, max_tokens=2000)  # 调用大模型API
        
        # 更新状态
        state.update({
            "expert_feedback": expert_feedback,
            "expert_review_complete": True
        })
        st.success("✅ 专家分析完成！")
        
        return state

    def recommendation_generation_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """推荐生成节点：调用外部服务获取相关论文推荐"""
        st.info("🌐 正在获取相关论文推荐...")

        try:
            keywords = state.get("recommendation_keywords", [])
            themes = state.get("recommendation_themes", [])

            if not keywords and not themes:
                st.warning(f"⚠️无法生成推荐：未提取到关键词或主题：{str(e)}")
                state.update({
                    "recommendations": "无法生成推荐",
                    "recommendation_complete": True
                })
                return state

            # 构造搜索查询
            search_query = " ".join(keywords + themes)

            # 模拟调用论文数据库（如Semantic Scholar、ArXiv等）
            # 实际应用中替换为真实API调用
            prompt = f"""请基于以下关键词，模拟从学术数据库中获取相关论文推荐：

关键词：{search_query}

请返回5篇相关论文，格式如下：
1. [论文标题] - [作者], [发表年份]
   - 领域：[研究领域]
   - 摘要：[简短摘要]

2. [论文标题] - [作者], [发表年份]
   - 领域：[研究领域]
   - 摘要：[简短摘要]

...
"""
            recommendations = self.call_glm(prompt, max_tokens=1000)

            # 更新状态
            state.update({
                "recommendations": recommendations,
                "recommendation_complete": True
            })

            st.success("✅ 论文推荐获取完成！")

        except Exception as e:
            st.error(f"❌ 获取推荐失败：{str(e)}")
            state.update({
                "recommendations": "推荐获取失败",
                "recommendation_complete": False,
                "error": str(e)
            })

        return state
